package com.cloud.common.elasticsearch;

import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.aggregations.Aggregation;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.bucket.terms.ParsedTerms;
import org.elasticsearch.search.aggregations.metrics.*;
import org.springframework.data.domain.PageRequest;

import java.util.*;

public class AggregationHelper {
    /**
     * 聚合分析的分页对象，用于控制不返回搜索的数据
     *
     * @return PageRequest
     */
    public static PageRequest aggregationPage() {
        return PageRequest.of(0, 1);
    }

    /**
     * 分组聚合分析时，同时返回命中数据的其他字段，所以返回数量1条就满足了
     *
     * @param fieldNames
     * @return TopHitsAggregationBuilder
     */
    public static TopHitsAggregationBuilder topHitsBuilder(String[] fieldNames) {
        return AggregationBuilders.topHits("fetch_field_hits")
                .fetchSource(fieldNames, null)
                .size(1);
    }

    /**
     * 解析聚合结果
     *
     * @param aggList
     * @return BucketResult
     */
    public static AggregationResult parse(List<Aggregation> aggList) {
        AggregationResult aggregationResult = new AggregationResult();
        aggList.forEach(aggregation -> {
            // 度量统计
            AggregationMetric aggMetric = getMetric(aggregation);
            if (aggMetric != null) {
                aggregationResult.addMetric(aggMetric);
                return;
            }
            // 分组统计
            if (aggregation instanceof ParsedTerms) {
                ParsedTerms terms = (ParsedTerms) aggregation;
                terms.getBuckets().forEach(bucket -> {
                    AggregationMetric metric = new AggregationMetric();
                    metric.setName(terms.getName());
                    metric.setKey(bucket.getKeyAsString());
                    metric.setCount(bucket.getDocCount());
                    aggregationResult.addMetric(metric);

                    bucket.getAggregations().asList().forEach(agg -> {
                        // 字段属性
                        if (agg instanceof ParsedTopHits) {
                            ParsedTopHits topHits = (ParsedTopHits) agg;
                            metric.setFieldJson(parseTopHits(topHits));
                            return;
                        }
                        // 分组统计
                        if (agg instanceof ParsedTerms) {
                            parseTermsAggregation(agg, metric);
                            return;
                        }
                        // 度量统计
                        AggregationMetric subAggMetric = getMetric(agg);
                        if (subAggMetric != null) {
                            metric.addSubMetric(subAggMetric);
                        }
                    });
                });
            }
        });
        return aggregationResult;
    }

    private static AggregationMetric getMetric(Aggregation aggregation) {
        // 累计计数
        if (aggregation instanceof ParsedValueCount) {
            ParsedValueCount parsedValueCount = (ParsedValueCount) aggregation;
            AggregationMetric metric = new AggregationMetric();
            metric.setName(parsedValueCount.getName());
            metric.setCount(parsedValueCount.getValue());
            return metric;
        }
        // 去重计数
        if (aggregation instanceof ParsedCardinality) {
            ParsedCardinality parsedCardinality = (ParsedCardinality) aggregation;
            AggregationMetric metric = new AggregationMetric();
            metric.setName(parsedCardinality.getName());
            metric.setCount(parsedCardinality.getValue());
            return metric;
        }
        // 汇总值
        if (aggregation instanceof ParsedSum) {
            ParsedSum parsedSum = (ParsedSum) aggregation;
            AggregationMetric metric = new AggregationMetric();
            metric.setName(parsedSum.getName());
            metric.setSum(parsedSum.getValue());
            return metric;
        }
        // 平均值
        if (aggregation instanceof ParsedAvg) {
            ParsedAvg parsedAvg = (ParsedAvg) aggregation;
            AggregationMetric metric = new AggregationMetric();
            metric.setName(parsedAvg.getName());
            metric.setAvg(parsedAvg.getValue());
            return metric;
        }
        // 最小值
        if (aggregation instanceof ParsedMin) {
            ParsedMin parsedMin = (ParsedMin) aggregation;
            AggregationMetric metric = new AggregationMetric();
            metric.setName(parsedMin.getName());
            metric.setMin(parsedMin.getValue());
            return metric;
        }
        // 最大值
        if (aggregation instanceof ParsedMax) {
            ParsedMax parsedMax = (ParsedMax) aggregation;
            AggregationMetric metric = new AggregationMetric();
            metric.setName(parsedMax.getName());
            metric.setMax(parsedMax.getValue());
            return metric;
        }
        return null;
    }

    private static void parseTermsAggregation(Aggregation aggregation, AggregationMetric aggregationMetric) {
        Map<AggregationMetric, List<Aggregation>> bucketMetricMap = new HashMap<>();
        ParsedTerms terms = (ParsedTerms) aggregation;

        // 遍历分组
        terms.getBuckets().forEach(bucket -> {
            AggregationMetric subAggMetric = new AggregationMetric();
            subAggMetric.setName(terms.getName());
            subAggMetric.setKey(bucket.getKeyAsString());
            subAggMetric.setCount(bucket.getDocCount());
            aggregationMetric.addSubMetric(subAggMetric);
            bucketMetricMap.put(subAggMetric, bucket.getAggregations().asList());
        });

        // 遍历下级
        bucketMetricMap.entrySet().forEach(bucketEntry -> {
            AggregationMetric metric = bucketEntry.getKey();
            List<Aggregation> aggList = bucketEntry.getValue();
            for (Aggregation subAgg : aggList) {
                // 递归处理
                if (subAgg instanceof ParsedTerms) {
                    parseTermsAggregation(subAgg, metric);
                    continue;
                }
                // 度量统计
                AggregationMetric subAggMetric = getMetric(subAgg);
                if (subAggMetric != null) {
                    metric.addSubMetric(subAggMetric);
                }
            }
        });
    }

    private static String parseTopHits(ParsedTopHits topHits) {
        SearchHit documentFields = Arrays.stream(topHits.getHits().getHits()).findFirst().orElse(null);
        return documentFields.getSourceAsString();
    }
}
